﻿#include "JZModelManager.h"
#include "JZNodeUtils.h"
#include "../JZModuleConfigFactory.h"

//JZModelConfigEnum
QDataStream& operator<<(QDataStream& s, const JZModelConfigEnum& config)
{
	JZModuleConfigFactory<JZModelConfig>::instance()->saveToStream(s,config);
	return s;
}

QDataStream& operator>>(QDataStream& s, JZModelConfigEnum& config)
{
	JZModuleConfigFactory<JZModelConfig>::instance()->loadFromStream(s, config);
	return s;
}

//JZModelManagerConfig
int JZModelManagerConfig::indexOfModel(QString name)
{
    for (int i = 0; i < modelList.size(); i++)
    {
        if (modelList[i]->name == name)
            return i;
    }
    return -1;
}

QDataStream& operator<<(QDataStream& s, const JZModelManagerConfig& config)
{
	s << config.modelList;
	return s;
}

QDataStream& operator>>(QDataStream& s, JZModelManagerConfig& config)
{
	s >> config.modelList;
	return s;
}

//JZModelManager
JZModelManager::JZModelManager(QObject* parent)
	:QObject(parent)
{
}

JZModelManager::~JZModelManager()
{
	qDeleteAll(m_models);
	m_models.clear();
}

void JZModelManager::setConfig(const JZModelManagerConfig& config)
{
	m_config = config;
    init();
}

JZModelManagerConfig JZModelManager::config()
{
	return m_config;
}

void JZModelManager::init()
{
    //clear
    qDeleteAll(m_models);
    m_models.clear();

    //init
	for (int i = 0; i < m_config.modelList.size(); i++)
	{
		JZModel* model = createModel(m_config.modelList[i]);
		m_models.push_back(model);
	}
}

JZModel* JZModelManager::model(QString name)
{
	for (int i = 0; i < m_config.modelList.size(); i++)
	{
		if(m_config.modelList[i]->name == name)
			return m_models[i];
	}
	return nullptr;
}

QList<JZModel*> JZModelManager::modelList()
{
    return m_models;
}
	
JZModel* JZModelManager::createModel(JZModelConfigEnum config)
{
	JZModel* model = nullptr;
	if (config->type == Model_Yolo)
	{
		JZYolo* yolo = new JZYolo();
		model = yolo;
	}
	else
	{
		Q_ASSERT(0);
	}

	model->setConfig(config);
	return model;
}

void JZModelInit(JZModelManager* inst, const QByteArray& buffer)
{
	JZModelManagerConfig config = JZNodeUtils::fromBuffer<JZModelManagerConfig>(buffer);
	inst->setConfig(config);
}

JZModel *JZModelGet(JZModelManager *inst, QString name)
{
    JZModel *model = inst->model(name);
    if (!model)
        throw std::runtime_error(qUtf8Printable("no model " + name));

    if (!model->isInit())
    {
        if(!model->init())
            throw std::runtime_error(qUtf8Printable("init model " + name + " failed"));
    }

    return model;
}

QList<JZYoloResult> JZYoloForward(JZModelManager *inst, QString name, cv::Mat mat)
{
    JZYolo *model = dynamic_cast<JZYolo*>(JZModelGet(inst, name));
    return model->forward(mat);
}